import csv
import gzip
import math
import time

import numpy as np
import matplotlib
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pack_padded_sequence

matplotlib.use('TkAgg')

trainFilename = "../data/Names/names_train.csv.gz"
testFilename = "../data/Names/names_test.csv.gz"


class NameDataset(Dataset):
    def __init__(self, is_train_set=True):
        filename = trainFilename if is_train_set else testFilename
        # 使用gzip.open()以文本模式打开压缩文件
        # 通过csv.reader()创建CSV文件读取器来解析压缩文件中的数据
        # r: 表示以只读模式打开文件 read
        # t: 表示以文本模式打开文件 text
        # rt 模式会将压缩文件解压并作为文本文件处理
        with gzip.open(filename, "rt") as f:
            # csv.reader(f) 返回的是一个迭代器，它逐行读取 CSV 文件内容
            # list(reader) 会遍历这个迭代器，将所有行数据收集到一个列表中
            # 每一行 CSV 数据都会作为列表的一个元素（通常是字符串列表）
            reader = csv.reader(f)
            # 将CSV文件读取器(reader)中的所有行数据转换为列表格式
            rows = list(reader)
        self.len = len(rows)  # 数据集长度
        self.names = [row[0] for row in rows]  # 姓名
        self.countries = [row[1] for row in rows]  # 国家
        self.country_list = list(sorted(set(self.countries)))  # 国家列表
        self.country_dict = self.getCountryDict()  # 国家字典
        self.country_num = len(self.country_list)  # 国家数量 -> 输出层维度（输出类别数）

    # getitem 方法的作用就是让 NameDataset 对象可以像列表一样使用索引访问。
    # 具体来说：
    # 实现索引访问：当使用 dataset[index] 时，会调用 getitem(self, index) 方法
    # 返回数据项：该方法返回指定索引位置的姓名和国家信息
    # 支持迭代操作：配合 len 方法，使数据集支持 for 循环遍历等操作
    # 符合PyTorch Dataset规范：这是 PyTorch Dataset 类要求实现的核心方法之一
    def __getitem__(self, index):
        return self.names[index], self.country_dict[self.countries[index]]

    def __len__(self):
        return self.len

    # 获取国家字典
    def getCountryDict(self):
        country_dict = dict()
        for idx, country_name in enumerate(self.country_list):
            country_dict[country_name] = idx
        return country_dict


trainSet = NameDataset(is_train_set=True)
testSet = NameDataset(is_train_set=False)

BATCH_SIZE = 64
# 批量处理数据：将数据集按指定的 batch_size 进行分批处理
# 数据打乱：通过 shuffle 参数控制是否随机打乱数据顺序
# 并行加载：支持多进程加载数据，提高训练效率
# 自动批处理：自动将单个样本组合成批次张量
trainLoader = DataLoader(trainSet, batch_size=BATCH_SIZE, shuffle=True)
testLoader = DataLoader(testSet, batch_size=BATCH_SIZE, shuffle=False)

EPOCHS = 18  # 训练轮数
CHARS_NUM = 128  # 输入字符数量
INPUT_SIZE = 128  # 输入向量维度
HIDDEN_SIZE = 128  # 隐藏层向量维度
OUTPUT_SIZE = trainSet.country_num  # 输出向量维度
NUM_LAYERS = 2  # 循环神经网络层数
USE_GPU = torch.cuda.is_available()  # 判断是否有可用的GPU设备


class RNNClassifier(nn.Module):
    def __init__(self,
                 char_num,
                 input_size,
                 hidden_size,
                 output_size,
                 num_layers=1,
                 batch_first=False,
                 bidirectional=False):
        super(RNNClassifier, self).__init__()
        self.n_directions = 2 if bidirectional else 1  # GRU 的单向与双向模式
        # char_num=128：词汇表大小，这里表示 ASCII 字符集大小
        # input_size=128：嵌入向量的维度
        self.embedding = nn.Embedding(char_num, input_size)  # 嵌入层
        # bidirectional 参数控制 GRU 是否使用双向模式，两者的区别如下：
        # bidirectional=False（单向GRU）：
        # 只从左到右处理序列数据
        # 每个时间步只考虑历史信息（前面的输入）
        # 输出维度：(batch_size, seq_len, hidden_size)
        # bidirectional=True（双向GRU）：
        # 同时从两个方向处理序列数据：
        # 前向：从左到右处理（同单向GRU）
        # 后向：从右到左处理
        # 每个时间步同时考虑历史和未来信息
        # 输出维度：(batch_size, seq_len, hidden_size*2)（因为有两个方向的隐藏状态拼接）
        self.gru = nn.GRU(input_size,
                          hidden_size,
                          num_layers,
                          batch_first=batch_first,
                          bidirectional=bidirectional)  # GRU层
        self.fc = nn.Linear(hidden_size * self.n_directions, output_size)  # 全连接层

    def forward(self, input, seq_lengths):
        input = input.t()  # (batch_size, seq_len) -> (seq_len, batch_size)
        embedding = self.embedding(input)
        # pack_padded_sequence 的作用：把一个按序列长度做了 padding 的批次，压缩成“只包含有效时间步”的稀疏表示，以便 RNN/LSTM/GRU 跳过 padding 部分 做计算与反传。
        # 通常把变长序列右侧补零（padding）凑成同样的 T_max；
        # pack_padded_sequence(emb, lengths, …) 会丢弃这些补零步，把张量打包成 PackedSequence；
        # RNN 读入 PackedSequence 后，只在每条序列真实长度上计算，更快、更准（不会让 padding 影响状态/梯度）。
        gru_input = pack_padded_sequence(embedding, seq_lengths)
        output, hidden = self.gru(gru_input)
        # PyTorch 规定“每一层先放 forward，再放 backward”，然后再进入下一层，所以最后一层的两个方向正好落在 -2（forward）与 -1（backward）。
        # 做分类/序列级表示时，常用 torch.cat([hidden[-2], hidden[-1]], dim=1) 把最后一层的双向信息拼在一起，得到 [B, 2*H] 的句子/序列表示。
        if self.n_directions == 2:
            hidden_cat = torch.cat([hidden[-2], hidden[-1]], dim=1)
        else:
            hidden_cat = hidden[-1]
        fc_output = self.fc(hidden_cat)
        return fc_output


def name2list(name: str):
    arr = [ord(c) for c in name]
    return arr, len(arr)


def create_tensor(tensor: torch.Tensor):
    if USE_GPU:
        return tensor.cuda()
    else:
        return tensor


# 把每个名字转为 ASCII 数组（ord(c)）。
# 在一个 batch 内需要 padding 到统一长度（右侧补 0）。
# 获取每个样本的 真实长度，并对 batch 做 按长度降序排序（pack_padded_sequence 的要求）。
# 返回三样：
# seq_tensor：形状 [B, S]（后续会转置成 [S, B] 喂给 RNN）
# seq_lengths：形状 [B] 的降序长度
# countries：形状 [B] 的类别索引
def make_tensors(names, countries):
    # names: List[str], countries: Tensor[int]（来自 DataLoader 的 batch）
    sequences_and_lengths = [name2list(name) for name in names]
    name_sequences = [sl[0] for sl in sequences_and_lengths]  # 获取每个名字的 ASCII 数组
    seq_lengths = torch.LongTensor([sl[1] for sl in sequences_and_lengths])  # 获取每个名字 ASCII 数组长度
    countries = torch.LongTensor(countries)

    # padding 到同一长度 -> [B, S]
    seq_tensor = torch.zeros(len(name_sequences), seq_lengths.max().item()).long()
    for idx, (seq, seq_len) in enumerate(zip(name_sequences, seq_lengths)):
        seq_tensor[idx, :seq_len] = torch.LongTensor(seq)

    # 按长度降序排序（pack_padded_sequence 要求）
    # 参数 0 表示沿着第0个维度（行维度）进行排序
    seq_lengths, perm_idx = seq_lengths.sort(0, descending=True)
    seq_tensor = seq_tensor[perm_idx]
    countries = countries[perm_idx]

    # 放到 GPU（若可用）
    return create_tensor(seq_tensor), seq_lengths, create_tensor(countries)


# s = time.time() - since - 计算当前时间与给定时间戳的差值（秒数）
# m = math.floor(s / 60) - 将总秒数转换为分钟数（向下取整）
# s -= m * 60 - 计算剩余的秒数
# return '%dm %ds' % (m, s) - 返回格式化的字符串，显示为"X分钟 Y秒"
def time_since(since):
    s = time.time() - since
    m = math.floor(s / 60)
    s -= m * 60
    return '%dm %ds' % (m, s)


def trainModel(classifier, criterion, optimizer, epoch, start, train_loader, train_set):
    classifier.train()
    train_loss = 0  # 训练损失
    seen = 0  # 已处理的样本数

    for i, (names, countries) in enumerate(train_loader, 1):
        seq_tensor, seq_lengths, countries = make_tensors(names, countries)
        outputs = classifier(seq_tensor, seq_lengths)
        loss = criterion(outputs, countries)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        bsz = seq_tensor.size(0)
        seen += bsz
        train_loss += loss.item() * bsz

        if i % 10 == 0:
            avg_loss = train_loss / seen
            print(f'[{time_since(start)}] Epoch {epoch} [{seen}/{len(train_set)}] loss={avg_loss:.6f}')

    return train_loss / max(1, seen)


# @torch.no_grad() 装饰器的作用与 with torch.no_grad(): 上下文管理器相同
# 使用装饰器后，被装饰的函数内部会自动禁用梯度计算
# 无需再在函数内部手动添加 with torch.no_grad(): 代码块
@torch.no_grad()
def testModel(classifier, test_loader, test_set):
    classifier.eval()
    correct = 0
    total = len(test_set)
    for i, (names, countries) in enumerate(test_loader, 1):
        seq_tensor, seq_lengths, countries = make_tensors(names, countries)
        outputs = classifier(seq_tensor, seq_lengths)
        pred = outputs.argmax(dim=1)
        correct += (pred == countries).sum().item()
    acc = correct / total if total > 0 else 0
    print(f'Test set: Accuracy {correct}/{total} = {acc:.4f}')
    return acc


def main():
    classifier = RNNClassifier(CHARS_NUM,
                               INPUT_SIZE,
                               HIDDEN_SIZE,
                               OUTPUT_SIZE,
                               NUM_LAYERS,
                               batch_first=False,
                               bidirectional=True)
    if USE_GPU:
        classifier = classifier.cuda()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(classifier.parameters(), lr=1e-3)

    start = time.time()
    acc_list = []
    for epoch in range(1, EPOCHS + 1):
        trainModel(classifier, criterion, optimizer, epoch, start, trainLoader, trainSet)
        acc = testModel(classifier, testLoader, testSet)
        acc_list.append(acc)

    epoch_idx = np.arange(1, len(acc_list) + 1, 1)
    plt.figure()
    plt.plot(epoch_idx, np.array(acc_list), marker='o')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.title('Test Accuracy vs. Epochs')
    plt.grid(True)
    plt.show()

    save_path = "../models/name_country_state_dict.pth"
    torch.save(classifier.state_dict(), save_path)
    print(f"模型参数已保存到: {save_path}")


if __name__ == '__main__':
    main()
